{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 求解一维传热"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\whn\\AppData\\Local\\Temp\\ipykernel_16396\\2730964139.py:61: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  x_b, t_b, u_b = map(lambda v: torch.tensor(v, dtype=torch.float32), (x_b, t_b, u_b))\n",
      "C:\\Users\\whn\\AppData\\Local\\Temp\\ipykernel_16396\\2730964139.py:62: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  x_i, t_i, u_i = map(lambda v: torch.tensor(v, dtype=torch.float32), (x_i, t_i, u_i))\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# 物理参数\n",
    "alpha = 0.01  # 热扩散系数\n",
    "\n",
    "# 定义 Transformer 作为特征提取器\n",
    "class TransformerFeatureExtractor(nn.Module):\n",
    "    def __init__(self, input_dim=2, embed_dim=32, num_heads=4, num_layers=2):\n",
    "        super(TransformerFeatureExtractor, self).__init__()\n",
    "        self.embedding = nn.Linear(input_dim, embed_dim)  # 线性变换到更高维度\n",
    "        encoder_layers = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=64)\n",
    "        self.transformer = nn.TransformerEncoder(encoder_layers, num_layers=num_layers)\n",
    "        self.fc_out = nn.Linear(embed_dim, 1)  # 最终输出 u(x,t)\n",
    "\n",
    "    def forward(self, x, t):\n",
    "        xt = torch.cat((x, t), dim=1)  # 连接 x 和 t\n",
    "        xt_embedded = self.embedding(xt).unsqueeze(1)  # (batch, seq_len=1, embed_dim)\n",
    "        xt_transformed = self.transformer(xt_embedded).squeeze(1)  # Transformer 处理\n",
    "        return self.fc_out(xt_transformed)\n",
    "\n",
    "# 计算物理损失\n",
    "def pde_loss(model, x, t):\n",
    "    x.requires_grad = True\n",
    "    t.requires_grad = True\n",
    "    u = model(x, t)\n",
    "\n",
    "    # 计算一阶导数 du/dt\n",
    "    u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]\n",
    "\n",
    "    # 计算二阶导数 d²u/dx²\n",
    "    u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]\n",
    "    u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]\n",
    "\n",
    "    # 物理方程残差\n",
    "    return torch.mean((u_t - alpha * u_xx) ** 2)\n",
    "\n",
    "# 生成训练数据\n",
    "N_f = 10000  # 物理方程点\n",
    "N_b = 100     # 边界点数\n",
    "N_i = 100     # 初始点数\n",
    "\n",
    "# 随机生成物理方程点 (x, t)\n",
    "x_f = torch.rand((N_f, 1), dtype=torch.float32, requires_grad=True)\n",
    "t_f = torch.rand((N_f, 1), dtype=torch.float32, requires_grad=True)\n",
    "\n",
    "# 边界条件数据 (u(0,t) = u(1,t) = 0)\n",
    "x_b = torch.cat([torch.zeros((N_b, 1)), torch.ones((N_b, 1))], dim=0)\n",
    "t_b = torch.rand((2 * N_b, 1))\n",
    "u_b = torch.zeros((2 * N_b, 1))\n",
    "\n",
    "# 初始条件数据 (u(x,0) = sin(pi*x))\n",
    "x_i = torch.rand((N_i, 1))\n",
    "t_i = torch.zeros((N_i, 1))\n",
    "u_i = torch.sin(np.pi * x_i)\n",
    "\n",
    "# 转换为 PyTorch 张量\n",
    "x_b, t_b, u_b = map(lambda v: torch.tensor(v, dtype=torch.float32), (x_b, t_b, u_b))\n",
    "x_i, t_i, u_i = map(lambda v: torch.tensor(v, dtype=torch.float32), (x_i, t_i, u_i))\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\Conda\\lib\\site-packages\\torch\\nn\\modules\\transformer.py:282: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.self_attn.batch_first was not True(use batch_first for better inference performance)\n",
      "  warnings.warn(f\"enable_nested_tensor is True, but self.use_nested_tensor is False because {why_not_sparsity_fast_path}\")\n"
     ]
    }
   ],
   "source": [
    "# 初始化模型和优化器\n",
    "model = TransformerFeatureExtractor()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "# 训练\n",
    "epochs = 5000\n",
    "for epoch in range(epochs):\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # 计算损失\n",
    "    loss_pde = pde_loss(model, x_f, t_f)\n",
    "    loss_bc = torch.mean((model(x_b, t_b) - u_b) ** 2)\n",
    "    loss_ic = torch.mean((model(x_i, t_i) - u_i) ** 2)\n",
    "    loss = loss_pde + loss_bc + loss_ic\n",
    "\n",
    "    # 反向传播和优化\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    # 打印训练进度\n",
    "    if epoch % 500 == 0:\n",
    "        print(f\"Epoch {epoch}, Loss: {loss.item()}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 预测结果\n",
    "x_test = np.linspace(0, 1, 100)\n",
    "t_test = np.linspace(0, 1, 100)\n",
    "X, T = np.meshgrid(x_test, t_test)\n",
    "X_test = torch.tensor(X.flatten(), dtype=torch.float32).view(-1, 1)\n",
    "T_test = torch.tensor(T.flatten(), dtype=torch.float32).view(-1, 1)\n",
    "\n",
    "with torch.no_grad():\n",
    "    U_pred = model(X_test, T_test).numpy().reshape(100, 100)\n",
    "\n",
    "# 绘制结果\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.imshow(U_pred, extent=[0, 1, 0, 1], origin=\"lower\", aspect=\"auto\", cmap=\"hot\")\n",
    "plt.colorbar(label=\"Temperature\")\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"t\")\n",
    "plt.title(\"Transformer-PINN Solution to 1D Heat Equation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
